package com.csust.zerocode.controller;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.csust.zerocode.ai.TestThinkingService;
import com.csust.zerocode.ai.TestThinkingServiceFactory;
import com.csust.zerocode.ai.model.message.*;
import com.csust.zerocode.common.BaseResponse;
import com.csust.zerocode.common.ResultUtils;
import com.csust.zerocode.exception.ErrorCode;
import com.csust.zerocode.exception.ThrowUtils;
import com.csust.zerocode.model.entity.User;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.PartialThinking;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.tool.ToolExecution;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;

/**
 * @ ClassName HealthController
 * @ Description TODO
 * @ Author 钰玟
 * @ Date 2025/9/9 下午9:21
 * Version 1.0
 **/
@RestController
@RequestMapping("/health")
public class HealthController {

    @Resource
    private TestThinkingService testThinkingService;

    @GetMapping("/")
    public BaseResponse<String> healthCheck() {
        return ResultUtils.success( "ok");
    }

    /**
     * 应用聊天生成代码（流式 SSE）
     * @param message 用户消息
     * @return 生成结果流
     */
    @GetMapping(value = "/test/think", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> testThinking(@RequestParam String message) {
        ThrowUtils.throwIf(StrUtil.isBlank(message), ErrorCode.PARAMS_ERROR, "用户消息不能为空");
        // 调用服务生成代码（流式）
        TokenStream tokenStream = testThinkingService.testThinking(message);
        // 转换为 ServerSentEvent 格式
        Flux<String> contentFlux = processTokenStream(tokenStream);
        return contentFlux
                .map(chunk -> {
                    String jsonData;
                    StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
                    StreamMessageTypeEnum enumByValue = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
                    if(enumByValue.equals(StreamMessageTypeEnum.AI_RESPONSE)){
                        AiResponseMessage aiResponseMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
                        String data = aiResponseMessage.getData();
                        // 将内容包装成JSON对象
                        Map<String, String> wrapper = Map.of("d", data);
                        jsonData = JSONUtil.toJsonStr(wrapper);
                        return ServerSentEvent.<String>builder()
                                .data(jsonData)
                                .event("message")
                                .build();
                    } else if (enumByValue.equals(StreamMessageTypeEnum.AI_REASONING)) {
                        AiReasoningMessage aiReasoningMessage = JSONUtil.toBean(chunk, AiReasoningMessage.class);
                        String data = aiReasoningMessage.getData();
                        // 将内容包装成JSON对象
                        Map<String, String> wrapper = Map.of("d", data);
                        jsonData = JSONUtil.toJsonStr(wrapper);
                        return ServerSentEvent.<String>builder()
                                .data(jsonData)
                                .event("reasoning")
                                .build();
                    }else{
                        return ServerSentEvent.<String>builder()
                                .data("")
                                .build();
                    }

                })
                .concatWith(Mono.just(
                        // 发送结束事件
                        ServerSentEvent.<String>builder()
                                .event("done")
                                .data("")
                                .build()
                ));
    }

    private Flux<String> processTokenStream(TokenStream tokenStream) {
        System.out.println(tokenStream);
        return Flux.create(sink->{
            tokenStream.onPartialResponse((String partialResponse) -> {
                        System.err.println("partialResponse: " + partialResponse);
                        AiResponseMessage aiResponseMessage = new AiResponseMessage(partialResponse);
                        sink.next(JSONUtil.toJsonStr(aiResponseMessage));
                    })
                    .onPartialThinking((PartialThinking partialThinking) -> {
                        String thinking = partialThinking.text();
//                        String wrapper = String.format("<think>")
                        AiReasoningMessage aiReasoningMessage = new AiReasoningMessage(thinking);
                        System.out.println("partialThinking: " + thinking);
                        sink.next(JSONUtil.toJsonStr(aiReasoningMessage));
                    })
                    .onPartialToolExecutionRequest((index, toolExecutionRequest) -> {
                        ToolRequestMessage toolRequestMessage = new ToolRequestMessage(toolExecutionRequest);
                        sink.next(JSONUtil.toJsonStr(toolRequestMessage));
                    })
                    .onToolExecuted((ToolExecution toolExecution) -> {
                        ToolExecutedMessage toolExecutedMessage = new ToolExecutedMessage(toolExecution);
                        sink.next(JSONUtil.toJsonStr(toolExecutedMessage));
                    })
                    .onCompleteResponse((ChatResponse response) -> {
                        sink.complete();
                    })
                    .onError((Throwable error) -> {
                        error.printStackTrace();
                        sink.error(error);
                    })
                    .start();
        });
    }
}
